use std::{collections::HashMap, sync::{Arc, atomic::{AtomicU64, Ordering}}, time::Duration};

use tokio::{sync::{mpsc::{UnboundedSender, UnboundedReceiver}, RwLock, oneshot}, io::{AsyncWriteExt, AsyncReadExt}, net::tcp::{OwnedWriteHalf, OwnedReadHalf}};

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub struct NodeId {
    pub id: u32,
    pub area_id: u32,
    pub zone_id: u32,
}

impl NodeId {
    pub async fn write_to<W: AsyncWriteExt + Unpin>(&self, writer: &mut W) -> tokio::io::Result<()> {
        writer.write_u32(self.id).await?;
        writer.write_u32(self.area_id).await?;
        writer.write_u32(self.zone_id).await?;

        Ok(())
    }

    pub async fn read_from<A: AsyncReadExt + Unpin>(reader: &mut A) -> tokio::io::Result<Self> {
        let id = reader.read_u32().await?;
        let area_id = reader.read_u32().await?;
        let zone_id = reader.read_u32().await?;

        Ok(NodeId {
            id,
            area_id,
            zone_id,
        })
    }
}

pub struct SubMessage {
    message_id: u32,
    data: Vec<u8>,
}

struct NodeMessage {
    from: NodeId,
    to: NodeId,
    channel: String,
    message: SubMessage,
}

impl NodeMessage {
    async fn write_to<W: AsyncWriteExt + Unpin>(&self, writer: &mut W) -> tokio::io::Result<()> {
        self.from.write_to(writer).await?;
        self.to.write_to(writer).await?;
        writer.write_u32(self.channel.len() as u32).await?;
        writer.write_all(self.channel.as_bytes()).await?;
        writer.write_u32(self.message.message_id).await?;
        writer.write_u32(self.message.data.len() as u32).await?;
        writer.write_all(&self.message.data).await?;

        Ok(())
    }

    async fn read_from<A: AsyncReadExt + Unpin>(reader: &mut A) -> tokio::io::Result<Self> {
        let from = NodeId::read_from(reader).await?;
        let to = NodeId::read_from(reader).await?;
        let channel_len = reader.read_u32().await? as usize;
        let mut channel = vec![0u8; channel_len];
        reader.read_exact(&mut channel).await?;
        let channel = String::from_utf8(channel).unwrap();
        let message_id = reader.read_u32().await?;
        let data_len = reader.read_u32().await? as usize;
        let mut data = vec![0u8; data_len];
        reader.read_exact(&mut data).await?;

        Ok(NodeMessage {
            from,
            to,
            channel,
            message: SubMessage {
                message_id,
                data,
            },
        })
    }
}

enum MessageSenderInner {
    Local(UnboundedSender<(NodeId, u32, Vec<u8>)>),
    Remote(NodeId, String, UnboundedSender<NodeMessage>),
    NoFind,
}

pub struct MessageSender {
    id: NodeId,
    inner: MessageSenderInner,
}

impl MessageSender {
    pub fn send(&self, message_id:u32, data: Vec<u8>) {
        match self.inner {
            MessageSenderInner::Local(ref tx) => {
                tx.send((self.id, message_id, data)).unwrap_or_default();
            },
            MessageSenderInner::Remote(ref id, ref channel, ref tx) => {
                tx.send(NodeMessage {
                    from: self.id,
                    to: id.clone(),
                    channel: channel.clone(),
                    message: SubMessage {
                        message_id,
                        data,
                    },
                }).unwrap_or_default();
            },
            MessageSenderInner::NoFind => {}
        }
    }
}

#[derive(Debug, Clone)]
enum GossipMessage {
    Route(NodeId, u32),
    Channel(String, NodeId),
}

impl GossipMessage {
    async fn write_to<W: AsyncWriteExt + Unpin>(&self, writer: &mut W) -> tokio::io::Result<()> {
        match self {
            GossipMessage::Route(id, unique_id) => {
                writer.write_u8(0).await?;
                id.write_to(writer).await?;
                writer.write_u32(*unique_id).await?;
            },
            GossipMessage::Channel(channel, id) => {
                writer.write_u8(1).await?;
                writer.write_u32(channel.len() as u32).await?;
                writer.write_all(channel.as_bytes()).await?;
                id.write_to(writer).await?;
            },
        }

        Ok(())
    }

    async fn read_from<A: AsyncReadExt + Unpin>(reader: &mut A) -> tokio::io::Result<Self> {
        let message_type = reader.read_u8().await?;
        match message_type {
            0 => {
                let id = NodeId::read_from(reader).await?;
                let unique_id = reader.read_u32().await?;

                Ok(GossipMessage::Route(id, unique_id))
            },
            1 => {
                let channel_len = reader.read_u32().await? as usize;
                let mut channel = vec![0u8; channel_len];
                reader.read_exact(&mut channel).await?;
                let channel = String::from_utf8(channel).unwrap();
                let id = NodeId::read_from(reader).await?;

                Ok(GossipMessage::Channel(channel, id))
            },
            _ => {
                Err(tokio::io::Error::new(tokio::io::ErrorKind::InvalidData, "invalid message type"))
            }
        }
    }
}

#[derive(Debug)]
struct Connection {
    unique_id: u64,
    close_tx: Option<oneshot::Sender<()>>,
    sender: UnboundedSender<NodeMessage>,
    gossip: UnboundedSender<GossipMessage>,
}

#[derive(Debug,Default)]
struct NodeInner {
    local_channels: HashMap<String, UnboundedSender<(NodeId, u32, Vec<u8>)>>,
    remote_channels: HashMap<String, Vec<NodeId>>,

    connections: HashMap<NodeId, Connection>,
    route_table: HashMap<NodeId, (u32, NodeId)>,

    gossip_wait_send_messages: Vec<GossipMessage>,
}

impl NodeInner {
    fn add_connection(&mut self, id: NodeId, unique_id: u64, close_tx: oneshot::Sender<()>, sender: UnboundedSender<NodeMessage>, gossip: UnboundedSender<GossipMessage>) {
        let close_tx = Some(close_tx);
        match self.connections.contains_key(&id) {
            true => {
                let conn = self.connections.get_mut(&id).unwrap();
                let old_close_tx = conn.close_tx.take().unwrap();
                old_close_tx.send(()).unwrap_or_default();

                conn.unique_id = unique_id;
                conn.close_tx = close_tx;
                conn.sender = sender;
                conn.gossip = gossip;
            },
            false => {
                self.connections.insert(id, Connection {
                    unique_id,
                    close_tx,
                    sender,
                    gossip,
                });
            }
        }
        
        self.add_route(id, id, 1);
    }

    fn remove_connection(&mut self, id: NodeId, unique_id: u64) {
        let remove = if let Some(conn) = self.connections.get(&id) {
            conn.unique_id == unique_id
        } else {
            false
        };

        if remove {
            self.connections.remove(&id);
        }
    }

    fn add_route(&mut self, from: NodeId, to: NodeId, cost: u32) {
        self.route_table.entry(from).and_modify(|f| {
            if f.0 > cost {
                f.0 = cost;
                f.1 = to;

                self.gossip_wait_send_messages.push(GossipMessage::Route(from, cost + 1));
            }
        }).or_insert_with(|| {
            self.gossip_wait_send_messages.push(GossipMessage::Route(from, cost + 1));
            (cost, to)
        });
    }

    fn add_local_channel(&mut self, channel: String) -> tokio::io::Result<UnboundedReceiver<(NodeId, u32, Vec<u8>)>> {
        if self.local_channels.contains_key(&channel) {
            return Err(tokio::io::Error::new(tokio::io::ErrorKind::AlreadyExists, "Channel already exists"));
        }

        let (tx, rx) = tokio::sync::mpsc::unbounded_channel();
        self.local_channels.insert(channel, tx);

        Ok(rx)
    }

    fn add_remote_channel(&mut self, channel: String, node: NodeId) {
        let remote = self.remote_channels.entry(channel.clone()).or_insert_with(Vec::new);
        if !remote.contains(&node) {
            remote.push(node);

            self.gossip_wait_send_messages.push(GossipMessage::Channel(channel, node));
        }
    }

    fn deliver_message(&self, id:&NodeId, message: NodeMessage) {
        if id == &message.to {
            if let Some(sender) = self.local_channels.get(message.channel.as_str()) {
                sender.send((message.from, message.message.message_id, message.message.data)).unwrap_or_default();
            }
        }
        else {
            let to_node = if let Some(to_node) = self.route_table.get(&message.to).map(|f| f.1) {
                to_node
            } else {
                return;
            };

            if let Some(conn) = self.connections.get(&to_node) {
                conn.sender.send(message).unwrap_or_default();
            }
        }
    }

    fn get_local_message_sender(&self, id: NodeId, channel: &str) -> Option<MessageSender> {
        if let Some(sender) = self.local_channels.get(channel) {
            return Some(MessageSender {
                id: id.clone(),
                inner: MessageSenderInner::Local(sender.clone()),
            });
        }

        None
    }

    fn get_remote_message_sender(&self, id: NodeId, channel: &str, node: Option<NodeId>) -> Option<MessageSender> {
        let nodes = if let Some(nodes) = self.remote_channels.get(channel) {
            if nodes.is_empty() {
                return None;
            }

            nodes
        } else {
            return None;
        };

        let node = if let Some(node) = node {
            if nodes.contains(&node) {
                node
            } else {
                return None;
            }
        } else {
            nodes.first().unwrap().clone()
        };

        let to_node = if let Some(to_node) = self.route_table.get(&node).map(|f| f.1) {
            to_node
        } else {
            return None;
        };

        if let Some(conn) = self.connections.get(&to_node) {
            return Some(MessageSender {
                id: id.clone(),
                inner: MessageSenderInner::Remote(node, channel.to_string(), conn.sender.clone()),
            });
        }

        None
    }

    fn get_message_sender(&self, id: NodeId, channel: &str, node: Option<NodeId>) -> MessageSender {
        let sender = if let Some(node) = node {
            if node == id {
                self.get_local_message_sender(id, channel)
            } else {
                self.get_remote_message_sender(id, channel, Some(node))
            }
        } else {
            self.get_local_message_sender(id, channel).or_else(|| self.get_remote_message_sender(id, channel, None))
        };

        sender.unwrap_or(MessageSender {
            id: id.clone(),
            inner: MessageSenderInner::NoFind,
        })
    }

    fn brocast_gossip(&mut self) {
        if self.gossip_wait_send_messages.is_empty() {
            return;
        }

        for conn in self.connections.values() {
            for gossip_message in &self.gossip_wait_send_messages {
                conn.gossip.send(gossip_message.clone()).unwrap_or_default();
            }
        }

        self.gossip_wait_send_messages.clear();
    }

    fn apply_gossip(&mut self, remote_id: NodeId, gossip_message: GossipMessage) {
        match gossip_message {
            GossipMessage::Route(node, cost) => {
                self.add_route(node, remote_id, cost);
            },
            GossipMessage::Channel(channel, node) => {
                self.add_remote_channel(channel, node);
            },
        }
    }

    fn get_init_gossip_messages(&self, id: &NodeId) -> Vec<GossipMessage> {
        let mut messages = Vec::new();

        for (from, (cost, _)) in &self.route_table {
            messages.push(GossipMessage::Route(from.clone(), cost + 1));
        }

        for (channel, _) in &self.local_channels {
            messages.push(GossipMessage::Channel(channel.clone(), id.clone()));
        }

        for (channel, nodes) in &self.remote_channels {
            for node in nodes {
                messages.push(GossipMessage::Channel(channel.clone(), node.clone()));
            }
        }

        messages
    }
    
}

const PING: u8 = 0;
const GOSSIP: u8 = 1;
const MESSAGE: u8 = 2;

fn setup_send_task(mut send_socket: OwnedWriteHalf, mut close_rx: oneshot::Receiver<()>, mut gossip_rx: UnboundedReceiver<GossipMessage>, mut message_rx: UnboundedReceiver<NodeMessage>) {
    tokio::spawn(async move {
        let mut instant = tokio::time::Instant::now();
        loop {
            instant = instant + Duration::from_secs(30);
            tokio::select! {
                _ = &mut close_rx => {
                    break;
                }
                _ = tokio::time::sleep_until(instant) => {
                    if let Err(_) = send_socket.write_u8(PING).await {

                    }
                }
                Some(message) = gossip_rx.recv() => {
                    if let Err(_) = send_socket.write_u8(GOSSIP).await {
                       
                    }

                    if let Err(_) = message.write_to(&mut send_socket).await {

                    }
                }
                Some(message) = message_rx.recv() => {
                    if let Err(_) = send_socket.write_u8(MESSAGE).await {
                    }

                    if let Err(_) = message.write_to(&mut send_socket).await {
                        
                    }
                }
            }
        }
    });
}

async fn poll_recv(mut recv_socket: OwnedReadHalf, id: &NodeId, remote_id: &NodeId, mut close_rx: oneshot::Receiver<()>, inner: &Arc<RwLock<NodeInner>>) -> tokio::io::Result<()>{
    loop {
        tokio::select! {
            _ = &mut close_rx => {
                break;
            }
            readed = tokio::time::timeout(Duration::from_secs(60), recv_socket.read_u8()) => {
                match readed {
                    Ok(message_type) => {
                        match message_type {
                            Ok(PING) => {
                                continue;
                            }
                            Ok(GOSSIP) => {
                                let message = GossipMessage::read_from(&mut recv_socket).await;
                                let message = match message {
                                    Ok(message) => message,
                                    Err(e) => {
                                        return Err(e);
                                    }
                                };

                                let mut inner_data = inner.write().await;
                                inner_data.apply_gossip(remote_id.clone(), message);
                                inner_data.brocast_gossip();
                                drop(inner_data);
                            }
                            Ok(MESSAGE) => {
                                let message = NodeMessage::read_from(&mut recv_socket).await;
                                let message = match message {
                                    Ok(message) => message,
                                    Err(e) => {
                                        return Err(e);
                                    }
                                };

                                let inner_data = inner.read().await;
                                inner_data.deliver_message(&id, message);
                                drop(inner_data);
                            }
                            Err(e) => {
                                return Err(e);
                            }
                            _ => {
                                return Err(tokio::io::Error::new(tokio::io::ErrorKind::InvalidData, "invalid message type"));
                            }
                        }
                    }
                    Err(e) => {
                        return Err(tokio::io::Error::new(tokio::io::ErrorKind::TimedOut, "recv message timeout"));
                    }
                }
            }
        }
    }

    Ok(())
}

async fn process_node_connection(mut socket: tokio::net::TcpStream, id: NodeId, unique_id: u64, inner: Arc<RwLock<NodeInner>>) -> tokio::io::Result<()> {
    id.write_to(&mut socket).await?;
    let remote_id = NodeId::read_from(&mut socket).await?;

    let inner_data = inner.read().await;
    let gossip_messages = inner_data.get_init_gossip_messages(&id);
    drop(inner_data);

    for message in gossip_messages.iter() {
        message.write_to(&mut socket).await?;
    }

    let (recv_socket, send_socket) = socket.into_split();
    let (close_sender_tx, close_sender_rx) = tokio::sync::oneshot::channel();
    let (gossip_tx, gossip_rx) = tokio::sync::mpsc::unbounded_channel();
    let (message_tx, message_rx) = tokio::sync::mpsc::unbounded_channel();

    setup_send_task(send_socket, close_sender_rx, gossip_rx, message_rx);

    let (close_tx, close_rx) = tokio::sync::oneshot::channel();
    let mut inner_data = inner.write().await;
    inner_data.add_connection(id, unique_id, close_tx, message_tx, gossip_tx);
    inner_data.brocast_gossip();
    drop(inner_data);

    poll_recv(recv_socket, &id, &remote_id, close_rx, &inner).await?;

    close_sender_tx.send(()).unwrap_or_default();

    let mut inner_data = inner.write().await;
    inner_data.remove_connection(id, unique_id);
    drop(inner_data);
    Ok(())
}

pub struct NodeNet {
    id: NodeId,
    next_unique_id: Arc<AtomicU64>,
    inner: Arc<RwLock<NodeInner>>,
}

impl NodeNet {
    pub fn new(id: NodeId) -> Self {
        Self {
            id,
            next_unique_id: Arc::new(AtomicU64::new(0)),
            inner: Arc::new(RwLock::new(NodeInner::default())),
        }
    }

    pub fn id(&self) -> &NodeId {
        &self.id
    }

    pub async fn get_sender(&self, channel: &str, node: Option<NodeId>) -> MessageSender {
        let inner = self.inner.read().await;
        inner.get_message_sender(self.id.clone(), channel, node)
    }

    pub async fn add_channel(&self, channel: &str) -> tokio::io::Result<UnboundedReceiver<(NodeId, u32, Vec<u8>)>> {
        let mut inner = self.inner.write().await;
        inner.add_local_channel(channel.to_string())
    }

    pub async fn listen(&self, addr: &str) -> tokio::io::Result<()> {
        let listener = tokio::net::TcpListener::bind(addr).await?;
        let inner = self.inner.clone();
        let id = self.id.clone();
        let next_unique_id = self.next_unique_id.clone();
        tokio::spawn(async move {
            loop {
                let (socket, _) = listener.accept().await.unwrap();
                let inner = inner.clone();
                let unique_id = next_unique_id.fetch_add(1, Ordering::Relaxed);
                tokio::spawn(async move {
                    if let Err(e) = process_node_connection(socket, id.clone(), unique_id, inner).await {
                        log::error!("process_node_connection error: {}", e);
                    }
                });
            }
        });

        Ok(())
    }

    pub async fn connect(&self, addr: &str) {
        let addr = addr.to_string();
        let inner = self.inner.clone();
        let id = self.id.clone();
        let unique_id = self.next_unique_id.fetch_add(1, Ordering::Relaxed);
        tokio::spawn(async move {
            loop {
                let socket = match tokio::net::TcpStream::connect(addr.clone()).await {
                    Ok(socket) => socket,
                    Err(_) => {
                        tokio::time::sleep(tokio::time::Duration::from_secs(1)).await;
                        continue;
                    }
                };

                if let Err(e) = process_node_connection(socket, id.clone(), unique_id, inner.clone()).await {
                    log::error!("process_node_connection error: {}", e);
                }
            }
        });
    }
}
